"""
框架1：中间突破+成分数据（组合预测）
"""

import pandas as pd
import numpy as np
from scipy.stats import gmean
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import TimeSeriesSplit
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.exponential_smoothing.ets import ETSModel
from statsmodels.tsa.forecasting.theta import ThetaModel
import logging
from concurrent.futures import ProcessPoolExecutor
import multiprocessing

logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s')

class Framework1:
    def __init__(self):
        """初始化分层预测框架"""
        self.regions = {
            "华北地区": ["北京市", "河北+天津", "山西省", "内蒙古自治区"],
            "东北地区": ["辽宁省", "吉林省", "黑龙江省"],
            "华东地区": ["上海市", "江苏省", "浙江省", "安徽省", "福建省", "江西省", "山东省"],
            "华中地区": ["河南省", "湖北省", "湖南省"],
            "华南地区": ["广东+海南", "广西壮族自治区"],
            "西南地区": ["四川+重庆", "贵州省", "云南省", "西藏自治区"],
            "西北地区": ["陕西省", "甘肃省", "青海省", "宁夏回族自治区", "新疆维吾尔自治区"]
        }
        self.logger = logging.getLogger(__name__)

    def load_data(self, file_path, sheet_name):
        """加载数据"""
        try:
            self.logger.info(f"正在从{file_path}加载{sheet_name}数据...")
            df = pd.read_excel(file_path, sheet_name=sheet_name)
            
            # 重命名列
            if "时间" in df.columns:
                df = df.rename(columns={"时间": "Year"})
            
            # 处理缺失值
            df = df.fillna(0)
            
            self.logger.info(f"数据加载成功，形状: {df.shape}")
            return df
        except Exception as e:
            self.logger.error(f"数据加载失败: {str(e)}")
            raise

    def create_hierarchical_structure(self, df):
        """创建层级结构"""
        try:
            self.logger.info("正在创建层级结构...")
            
            # 转换为长格式
            df_long = df.melt(id_vars=["Year"], 
                            var_name="Province", 
                            value_name="Grain")
            
            # 确保Grain列为数值类型
            df_long["Grain"] = pd.to_numeric(df_long["Grain"], errors="coerce")
            
            # 添加区域信息
            df_long["Area"] = df_long["Province"].map(
                {province: area
                 for area, provinces in self.regions.items()
                 for province in provinces}
            )
            
            # 计算区域和全国总量
            df_regions = df_long.groupby(["Year", "Area"])["Grain"].sum().reset_index()
            df_national = df_long.groupby("Year")["Grain"].sum().reset_index()
            
            self.logger.info("层级结构创建成功")
            return df_long, df_regions, df_national
        except Exception as e:
            self.logger.error(f"层级结构创建失败: {str(e)}")
            raise

    def prepare_ilr_data(self, df_long, df_regions, df_national):
        """准备ILR变换的数据"""
        try:
            self.logger.info("正在准备ILR变换数据...")
            
            # 计算区域占比
            df_regions = df_regions.merge(df_national, on="Year", suffixes=("", "_national"))
            df_regions["Proportion"] = df_regions["Grain"] / df_regions["Grain_national"]
            
            # 计算省份在区域中的占比
            df_provinces = df_long.merge(df_regions[["Year", "Area", "Grain"]], 
                                       on=["Year", "Area"], 
                                       suffixes=("", "_regional"))
            df_provinces["Proportion"] = df_provinces["Grain"] / df_provinces["Grain_regional"]
            
            # 转换为矩阵形式
            region_matrix = df_regions.pivot(index="Year", columns="Area", values="Proportion").values
            province_matrices = {}
            for area in self.regions:
                area_data = df_provinces[df_provinces["Area"] == area]
                if not area_data.empty:
                    province_matrices[area] = area_data.pivot(
                        index="Year", 
                        columns="Province", 
                        values="Proportion"
                    ).values
            
            self.logger.info("ILR变换数据准备完成")
            return region_matrix, province_matrices
        except Exception as e:
            self.logger.error(f"准备ILR变换数据失败: {str(e)}")
            raise

    def ilr_transform(self, X, epsilon=1e-10):
        """改进的ILR变换"""
        try:
            self.logger.info("正在执行ILR变换...")
            
            if isinstance(X, (pd.DataFrame, pd.Series)):
                X = X.values
                
            if not isinstance(X, np.ndarray):
                raise ValueError("输入数据必须是numpy数组或pandas对象")
                
            # 确保数据为浮点型
            X = X.astype(float)
            
            # 处理零值和极小值
            X = np.where(X == 0, epsilon, X)
            X = np.where(X < epsilon, epsilon, X)
            
            # 归一化
            X = X / X.sum(axis=1)[:, np.newaxis]
            
            # 计算基础矩阵
            n = X.shape[1]
            H = np.zeros((n-1, n))
            for i in range(n-1):
                H[i, :i+1] = 1/np.sqrt((i+1)*(i+2))
                H[i, i+1] = -np.sqrt((i+1)/(i+2))
            
            # 计算几何平均数并处理极小值
            g = gmean(X, axis=1)
            g = np.where(g < epsilon, epsilon, g)
            
            # 执行变换
            X_normalized = X / g[:, np.newaxis]
            X_normalized = np.where(X_normalized < epsilon, epsilon, X_normalized)
            
            # 安全的对数变换
            Z = np.log(X_normalized)
            
            # 应用Helmert矩阵
            result = np.dot(Z, H.T)
            
            self.logger.info("ILR变换完成")
            return result
        except Exception as e:
            self.logger.error(f"ILR变换失败: {str(e)}")
            raise

    def ilr_inverse(self, Y, n_components, epsilon=1e-10):
        """ILR逆变换"""
        try:
            self.logger.info("正在执行ILR逆变换...")
            
            # 构建Helmert矩阵
            H = np.zeros((Y.shape[1], n_components))
            for i in range(Y.shape[1]):
                H[i, :i+1] = 1/np.sqrt((i+1)*(i+2))
                H[i, i+1] = -np.sqrt((i+1)/(i+2))
            
            # 执行逆变换
            X = np.exp(np.dot(Y, H))
            
            # 处理数值问题
            X = np.where(X < epsilon, epsilon, X)
            
            # 归一化
            X = X / X.sum(axis=1)[:, np.newaxis]
            
            self.logger.info("ILR逆变换完成")
            return X
        except Exception as e:
            self.logger.error(f"ILR逆变换失败: {str(e)}")
            raise

    def prepare_time_series(self, data, freq='YE'):
        """准备时间序列数据"""
        try:
            # 确保数据是DataFrame
            if isinstance(data, np.ndarray):
                data = pd.DataFrame(data)
            
            # 创建时间索引
            years = data.index.astype(int)
            dates = pd.date_range(start=f"{min(years)}-12-31", 
                                end=f"{max(years)}-12-31", 
                                freq=freq)
            
            # 设置时间索引
            data.index = dates
            
            return data
            
        except Exception as e:
            self.logger.error(f"准备时间序列数据失败: {str(e)}")
            raise

    def fit_forecast_models(self, data, h=1):
        """拟合和预测时间序列模型"""
        try:
            self.logger.info("正在拟合预测模型...")
            
            # 准备时间序列数据
            ts_data = self.prepare_time_series(data)
            
            models = {
                'arima': lambda y: ARIMA(y, order=(1,1,1)).fit(),
                'ets': lambda y: ETSModel(y).fit(),
                'theta': lambda y: ThetaModel(y, period=1).fit()
            }
            
            # 使用时间序列交叉验证选择最佳模型
            tscv = TimeSeriesSplit(n_splits=5)
            best_models = {}
            
            for col in ts_data.columns:
                model_scores = {}
                for model_name, model_func in models.items():
                    scores = []
                    for train_idx, test_idx in tscv.split(ts_data[col]):
                        try:
                            train = ts_data[col].iloc[train_idx]
                            test = ts_data[col].iloc[test_idx]
                            model = model_func(train)
                            pred = model.forecast(len(test))
                            score = np.sqrt(np.mean((test - pred) ** 2))
                            scores.append(score)
                        except Exception as e:
                            self.logger.warning(f"模型 {model_name} 在列 {col} 上拟合失败: {str(e)}")
                            continue
                    if scores:
                        model_scores[model_name] = np.mean(scores)
                
                # 选择最佳模型
                if model_scores:
                    best_model_name = min(model_scores, key=model_scores.get)
                    self.logger.info(f"列 {col} 选择模型 {best_model_name}")
                    best_model = models[best_model_name](ts_data[col])
                    best_models[col] = best_model.forecast(h)
            
            return pd.DataFrame(best_models)
            
        except Exception as e:
            self.logger.error(f"模型拟合和预测失败: {str(e)}")
            raise

    def calculate_metrics(self, actual, predicted, train_actual=None):
        """计算评估指标"""
        try:
            self.logger.info("正在计算评估指标...")
            metrics = {}
            
            # 确保数据类型
            actual = np.asarray(actual, dtype=float)
            predicted = np.asarray(predicted, dtype=float)
            
            # 处理缺失值和异常值
            mask = ~(np.isnan(actual) | np.isnan(predicted))
            actual = actual[mask]
            predicted = predicted[mask]
            
            # RMSE
            metrics['rmse'] = float(np.sqrt(np.mean((actual - predicted) ** 2)))
            
            # MAPE
            non_zero_mask = actual != 0
            if np.any(non_zero_mask):
                metrics['mape'] = float(np.mean(np.abs((actual[non_zero_mask] - predicted[non_zero_mask]) / 
                                               actual[non_zero_mask])) * 100)
            else:
                metrics['mape'] = np.nan
            
            # MASE
            if train_actual is not None:
                train_actual = np.asarray(train_actual, dtype=float)
                naive_errors = np.abs(np.diff(train_actual))
                naive_mae = np.mean(naive_errors)
                if naive_mae != 0:
                    mae = np.mean(np.abs(actual - predicted))
                    metrics['mase'] = float(mae / naive_mae)
                else:
                    metrics['mase'] = np.nan
            
            # R²
            if len(actual) > 1:
                metrics['r2'] = float(1 - (np.sum((actual - predicted) ** 2) / 
                                   np.sum((actual - np.mean(actual)) ** 2)))
            
            self.logger.info(f"评估指标计算完成: {metrics}")
            return metrics
            
        except Exception as e:
            self.logger.error(f"评估指标计算失败: {str(e)}")
            raise

    def run_forecast(self, file_path, sheet_name, forecast_horizon=1):
        """运行完整的预测流程"""
        try:
            self.logger.info("开始运行预测流程...")
            
            # 1. 加载数据
            df = self.load_data(file_path, sheet_name)
            
            # 2. 创建分层结构
            df_long, df_regions, df_national = self.create_hierarchical_structure(df)
            
            # 3. 准备ILR变换数据
            region_matrix, province_matrices = self.prepare_ilr_data(df_long, df_regions, df_national)
            
            # 4. 执行ILR变换
            region_transformed = self.ilr_transform(region_matrix)
            province_transformed = {area: self.ilr_transform(matrix) 
                                 for area, matrix in province_matrices.items()}
            
            # 5. 准备时间序列数据
            region_ts = pd.DataFrame(region_transformed, 
                                   index=df_regions["Year"].unique())
            province_ts = {area: pd.DataFrame(transformed, 
                                            index=df_regions["Year"].unique())
                         for area, transformed in province_transformed.items()}
            
            # 6. 拟合预测模型
            region_forecasts = self.fit_forecast_models(region_ts, h=forecast_horizon)
            province_forecasts = {area: self.fit_forecast_models(ts_data, h=forecast_horizon)
                                for area, ts_data in province_ts.items()}
            
            # 7. ILR逆变换
            region_reconstructed = self.ilr_inverse(region_forecasts.values, 
                                                  region_matrix.shape[1])
            province_reconstructed = {area: self.ilr_inverse(forecasts.values,
                                                           province_matrices[area].shape[1])
                                    for area, forecasts in province_forecasts.items()}
            
            # 8. 计算评估指标
            metrics = {
                'region': self.calculate_metrics(
                    region_matrix[-forecast_horizon:], 
                    region_reconstructed
                ),
                'province': {
                    area: self.calculate_metrics(
                        province_matrices[area][-forecast_horizon:],
                        reconstructed
                    )
                    for area, reconstructed in province_reconstructed.items()
                }
            }
            
            self.logger.info("预测流程完成")
            return {
                'region_forecasts': region_reconstructed,
                'province_forecasts': province_reconstructed,
                'metrics': metrics
            }
        
        except Exception as e:
            self.logger.error(f"预测流程失败: {str(e)}")
            raise 